from langchain.agents import tool
from langchain_community.utilities import SerpAPIWrapper
from langchain_community.chat_models import ChatTongyi
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.utilities import SQLDatabase
from qdrant_client import QdrantClient
from langchain_qdrant import QdrantVectorStore
import os
import requests
from dotenv import load_dotenv
load_dotenv()
os.environ["SERPAPI_API_KEY"] = os.getenv("SERPAPI_API_KEY")

@tool
def serp_search(query: str) -> str:
    """只有需要了解实时信息或不知道的事情的时候才会使用这个工具."""
    serp = SerpAPIWrapper()
    return serp.run(query)

def query_from_db(query: str) -> str:
    """只有需要查询客户信息时才会使用这个工具。请根据用户的输入，获取到关键字如: 名字或者邮箱"""
    db = SQLDatabase.from_uri(f"mysql+pymysql://{os.getenv("db_user")}:{os.getenv("db_password")}@{os.getenv("db_host")}/{os.getenv("db_name")}")
    print(db.dialect)
    print(db.get_usable_table_names())
    cursor = db.run()
    return

if __name__ == '__main__':
    query_from_db("111")

# @tool
# def get_info_from_local(query: str) -> str:
#     """只有回答与2024年运势或龙年运势相关问题的时候会使用这个工具."""
#     client = Qdrant(
#         QdrantClient(path="/ai/local_qdrant"),
#         "local_documents_demo",
#         DashScopeEmbeddings(model="text-embedding-v2"),
#     )
#     retriever = client.as_retriever(search_type="mmr")
#     result = retriever.get_relevant_documents(query, collection_name="local_documents")
#     print("*****************", result)
#     return result

@tool
def get_info_from_local(query: str) -> str:
    """只有回答与2024年运势或龙年运势相关问题的时候会使用这个工具."""
    client = QdrantClient(url="http://127.0.0.1", port=6333, grpc_port=6334)
    retriever_qr = QdrantVectorStore(client, "local_documents_demo", DashScopeEmbeddings(model="text-embedding-v2"))
    retriever = retriever_qr.as_retriever(search_type="mmr")
    result = retriever.get_relevant_documents(query)
    return result

@tool
def bazi_cesuan(query: str):
    """只有做八字排盘或者八字测算的时候才会使用这个工具，需要输入用户姓名和出生年月日时，如果缺少用户姓名和出生年月日时则不使用"""
    url = f"https://api.yuanfenju.com/index.php/v1/Bazi/cesuan"
    prompt = ChatPromptTemplate.from_template(
        """你是一个参数查询助手，根据用户输入内容找出相关的参数并按json格式返回。
        JSON字段如下: -"api_key":"B57HzB2y9L2tLvd82Y04Qfuqy", -"name": "姓名", -"sex":"性别，0表示男 1表示女，根据姓名判断", 
        -"type":"日历类型，0农历，1公历，默认1", -"year": "出生年份 例：1996"， -"month":"出生月份，例8",
        -"day":"出生日期，例8", -"hours":"出生小时 例 14", -"minute":"0"， 如果没有找到相关参数，则需要提醒用户
        告诉你这些内容，只返回数据结构，不要有其他的评论， 用户输入:{query}。"""
    )
    parser = JsonOutputParser()
    prompt = prompt.partial(format_instructions=parser.get_format_instructions())
    print("bazi_cesuan prompt:", prompt)
    chain = prompt | ChatTongyi(model="qwen-turbo", temperature=0) | parser
    data = chain.invoke({"query": query})
    print("八字测算结果查询:", data)
    result = requests.post(url, data=data)
    if result.status_code == 200:
        print("===返回数据===")
        print(result.json())
        try:
            json = result.json()
            retrunstr = f"八字为: {json["data"]["bazi_info"]["bazi"]}";
            print(retrunstr)
            return retrunstr
        except Exception as e:
            return "八字查询失败, 可能是你忘记询问用户姓名或者出生年月日了。"
    else:
        return "技术错误，请告诉用户稍后再试。"

@tool
def jiemeng(query: str):
    """只有帮助用户解梦的时候才会使用这个工具，需要输入梦境内容."""
    url = f"https://api.yuanfenju.com/index.php/v1/Gongju/zhougong"
    LLM = ChatTongyi(
            model="qwen-turbo",
            temperature=0, 
            streaming=True
        )
    prompt = ChatPromptTemplate.from_template("根据输入的内容提取1个关键词，只返回关键词，输入内容:{topic}")
    prompt_value = prompt.invoke({"topic": query})
    keyword = LLM.invoke(prompt_value)
    print("=======提交数据=======")
    print(keyword.content)
    # result = requests.post(url, data={"api_key": os.getenv("API_KEY"), "title_zhougong": keyword.content})
    result = requests.post(url, data={"api_key": 'B57HzB2y9L2tLvd82Y04Qfuqy', "title_zhougong": keyword.content})
    if result.status_code == 200:
        print("=======返回数据=======")
        return result.text
    else:
        return "用户输入的信息缺失，提醒用户补充信息！"